import os
import json
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from config import Config
from openai import OpenAI


class RAGService:
    """知识抽取服务（含向量检索）"""
    
    def __init__(self):
        self.config = Config()
        self.embedding_model = None
        self.llm_client = None
        self._initialize()
    
    def _initialize(self):
        """初始化服务"""
        # 初始化嵌入模型（用于向量生成和检索）
        print("正在加载向量检索模型...")
        self.embedding_model = SentenceTransformer(
            self.config.EMBEDDING_MODEL,
            device='cpu',
            trust_remote_code=True
        )
        
        # 初始化LLM客户端
        self.llm_client = OpenAI(
            api_key=self.config.DASHSCOPE_API_KEY,
            base_url=self.config.DASHSCOPE_BASE_URL
        )
        
        print("✓ 知识抽取服务已就绪")
    
    def generate_embedding(self, text):
        """生成文本的向量嵌入"""
        if not text or not text.strip():
            return None
        embedding = self.embedding_model.encode([text])[0]
        # 转换为列表以便JSON序列化
        return embedding.tolist()
    
    def retrieve_relevant_context(self, query, dictionary_terms, document_chunks, top_k=5):
        """
        从领域词典和文档中检索相关上下文
        
        Args:
            query: 用户查询
            dictionary_terms: 领域词典列表 [{'id': xx, 'term': xx, 'definition': xx, 'embedding': json_str}, ...]
            document_chunks: 文档分块列表 [{'text': xx, 'embedding': json_str}, ...]
            top_k: 返回最相关的k个结果
        
        Returns:
            {'dictionary': [...], 'document': [...]}
        """
        results = {'dictionary': [], 'document': []}
        
        # 生成查询向量
        query_embedding = self.embedding_model.encode([query])[0]
        
        # 1. 从领域词典检索
        if dictionary_terms:
            dict_embeddings = []
            dict_items = []
            
            for term in dictionary_terms:
                if term.get('embedding'):
                    try:
                        emb = json.loads(term['embedding'])
                        dict_embeddings.append(emb)
                        dict_items.append(term)
                    except:
                        continue
            
            if dict_embeddings:
                # 创建临时Faiss索引
                dict_embeddings = np.array(dict_embeddings, dtype=np.float32)
                dimension = dict_embeddings.shape[1]
                index = faiss.IndexFlatL2(dimension)
                index.add(dict_embeddings)
                
                # 检索
                k = min(top_k, len(dict_items))
                distances, indices = index.search(
                    np.array([query_embedding], dtype=np.float32), k
                )
                
                for idx, distance in zip(indices[0], distances[0]):
                    if idx < len(dict_items):
                        results['dictionary'].append({
                            'item': dict_items[idx],
                            'distance': float(distance)
                        })
        
        # 2. 从文档分块检索
        if document_chunks:
            doc_embeddings = []
            doc_items = []
            
            for chunk in document_chunks:
                if chunk.get('embedding'):
                    try:
                        emb = json.loads(chunk['embedding'])
                        doc_embeddings.append(emb)
                        doc_items.append(chunk)
                    except:
                        continue
            
            if doc_embeddings:
                # 创建临时Faiss索引
                doc_embeddings = np.array(doc_embeddings, dtype=np.float32)
                dimension = doc_embeddings.shape[1]
                index = faiss.IndexFlatL2(dimension)
                index.add(doc_embeddings)
                
                # 检索
                k = min(top_k, len(doc_items))
                distances, indices = index.search(
                    np.array([query_embedding], dtype=np.float32), k
                )
                
                for idx, distance in zip(indices[0], distances[0]):
                    if idx < len(doc_items):
                        results['document'].append({
                            'item': doc_items[idx],
                            'distance': float(distance)
                        })
        
        return results
    
    def extract_entities(self, prompt, document_content, retrieved_context=None):
        """使用用户自定义提示词抽取知识（基于向量检索增强）"""
        
        # 构建领域知识参考
        dict_knowledge = ""
        if retrieved_context and retrieved_context.get('dictionary'):
            dict_items = retrieved_context['dictionary']
            if dict_items:
                dict_knowledge = "\n".join([
                    f"- {item['item']['term']}: {item['item']['definition']}"
                    for item in dict_items
                ])
        
        # 构建文档上下文（如果有检索到的相关片段）
        doc_context = document_content
        if retrieved_context and retrieved_context.get('document'):
            doc_items = retrieved_context['document']
            if doc_items:
                doc_context = "\n\n".join([item['item']['text'] for item in doc_items])
        
        enhanced_prompt = f"""你是一个专业的知识抽取专家。请从文档中抽取结构化知识。

【领域词典参考】（基于向量检索的相关专业术语）
{dict_knowledge if dict_knowledge else "无"}

【文档内容】（基于向量检索的相关片段）
{doc_context}

【用户需求】
{prompt}

【标准实体类型】（必须严格使用以下类型名称）
1. 设备名 - 设备的名称（如锅炉、汽轮机、水冷壁管等）
2. 故障状态 - 故障的具体表现（如爆管、振动超标、泄漏等）
3. 性能表征 - 性能参数变化（如压力下降、温度升高、流量减少等）
4. 故障原因 - 导致故障的原因（如结垢、腐蚀、磨损、老化等）
5. 解决办法 - 处理措施（如更换部件、调整参数、清洗等）

【抽取规则】
1. 实体类型名称必须使用上述5种标准类型之一
2. 实体值要具体、完整，避免抽象概念
3. 关系要准确反映实体之间的联系
4. 结合领域词典理解专业术语
5. 提取文档中明确提到的和合理推断的内容

【输出格式】
严格按照JSON格式返回，实体类型必须使用标准名称：
{{
    "entities": {{
        "设备名": ["具体设备1", "具体设备2"],
        "故障状态": ["具体故障1"],
        "性能表征": ["具体表现1"],
        "故障原因": ["具体原因1"],
        "解决办法": ["具体办法1"]
    }},
    "relations": [
        {{"subject": "实体A", "predicate": "关系名", "object": "实体B"}}
    ]
}}

【标准示例】
用户问：锅炉管爆了，想知道原因和处理方法
正确抽取：
{{
    "entities": {{
        "设备名": ["锅炉", "水冷壁管"],
        "故障状态": ["爆管", "过热爆裂"],
        "性能表征": ["主蒸汽压力下降", "温度异常升高"],
        "故障原因": ["管内结垢", "局部过热", "材质老化"],
        "解决办法": ["停炉检修", "更换爆管段", "清理水垢", "加强水质管理"]
    }},
    "relations": [
        {{"subject": "锅炉", "predicate": "发生故障", "object": "爆管"}},
        {{"subject": "爆管", "predicate": "表现", "object": "主蒸汽压力下降"}},
        {{"subject": "爆管", "predicate": "由原因引起", "object": "管内结垢"}},
        {{"subject": "管内结垢", "predicate": "导致", "object": "局部过热"}},
        {{"subject": "爆管", "predicate": "通过办法解决", "object": "更换爆管段"}}
    ]
}}

错误示例（过于抽象）：
{{
    "entities": {{}},
    "relations": [
        {{"subject": "原因", "predicate": "导致", "object": "爆管事故"}}
    ]
}}

现在开始抽取："""
        
        # 调用大模型
        try:
            completion = self.llm_client.chat.completions.create(
                model=self.config.QWEN_MODEL,
                messages=[
                    {"role": "system", "content": "你是一个专业的知识抽取专家，善于从文档中提取完整、具体的结构化知识，并能基于领域知识进行合理推断。"},
                    {"role": "user", "content": enhanced_prompt}
                ],
                temperature=0.3,
                stream=False
            )
            
            response_text = completion.choices[0].message.content.strip()
            
            # 解析JSON结果
            if "```json" in response_text:
                json_start = response_text.find("```json") + 7
                json_end = response_text.find("```", json_start)
                response_text = response_text[json_start:json_end].strip()
            elif "```" in response_text:
                json_start = response_text.find("```") + 3
                json_end = response_text.find("```", json_start)
                response_text = response_text[json_start:json_end].strip()
            
            result = json.loads(response_text)
            
            return {
                'success': True,
                'data': result
            }
            
        except json.JSONDecodeError as e:
            print(f"JSON解析错误: {e}")
            print(f"原始响应: {response_text}")
            return {
                'success': False,
                'message': f'JSON解析失败: {str(e)}',
                'raw_response': response_text
            }
        except Exception as e:
            print(f"抽取错误: {e}")
            return {
                'success': False,
                'message': f'知识抽取失败: {str(e)}'
            }
    
    def optimize_prompt(self, user_prompt, dictionary_knowledge=""):
        """优化用户提示词"""
        optimization_prompt = f"""你是一个提示词优化专家。请优化用户的知识抽取需求。

【用户原始需求】
{user_prompt}

【优化任务】
将用户的简单需求扩展为更清晰、更详细的描述：
1. 明确说明要抽取哪些类型的信息
2. 说明这些信息之间的关系
3. 使用自然语言表达，无需遵循固定格式
4. 保持简洁明了，避免冗长

【示例1】
用户输入：锅炉管爆了，想知道原因和处理方法
优化后：请从文档中抽取锅炉爆管的相关信息，包括涉及的设备名称、具体的故障现象、导致故障的原因、以及相应的处理措施和解决方案。

【示例2】
用户输入：抽取设备和故障
优化后：从文档中识别出现的设备名称、发生的故障类型、故障的具体表现、故障原因，以及设备与故障之间的关联关系。

【示例3】
用户输入：温度、压力
优化后：提取文档中提到的温度参数和压力参数，包括具体的数值、测量位置、异常变化情况，以及这些参数与设备状态的关系。

【注意】
- 直接输出优化后的提示词
- 使用自然流畅的语言
- 不要使用JSON、markdown等特殊格式
- 不要添加"以下是优化后的"等额外说明
"""
        
        try:
            completion = self.llm_client.chat.completions.create(
                model=self.config.QWEN_MODEL,
                messages=[
                    {"role": "system", "content": "你是一个提示词优化专家，擅长将用户的简单需求扩展为清晰、自然的描述。"},
                    {"role": "user", "content": optimization_prompt}
                ],
                temperature=0.5,
                stream=False
            )
            
            optimized_prompt = completion.choices[0].message.content.strip()
            
            return {
                'success': True,
                'optimized_prompt': optimized_prompt
            }
            
        except Exception as e:
            print(f"提示词优化错误: {e}")
            return {
                'success': False,
                'message': f'提示词优化失败: {str(e)}'
            }
    
    def generate_response(self, query, context='', history=None):
        """
        生成对话回复
        
        Args:
            query: 用户问题
            context: 检索到的上下文
            history: 历史对话记录列表，格式为 [{'role': 'user', 'content': '...'}, ...]
        
        Returns:
            str: AI回复文本
        """
        if history is None:
            history = []
        
        # 构建系统提示
        system_prompt = """你是一个专业的AI助手，擅长回答各种问题。你的任务是：
1. 基于提供的上下文信息回答用户问题
2. 如果上下文中没有相关信息，可以使用你的知识来回答
3. 保持回答简洁、准确、有帮助
4. 如果不确定答案，请诚实说明"""
        
        # 构建消息列表
        messages = [{"role": "system", "content": system_prompt}]
        
        # 添加历史对话（最多保留最近5轮）
        for hist_msg in history[-10:]:
            messages.append({
                "role": hist_msg['role'],
                "content": hist_msg['content']
            })
        
        # 构建当前问题（包含上下文）
        if context:
            user_content = f"""参考以下上下文信息：

{context}

用户问题：{query}

请基于上下文信息回答用户问题。如果上下文中没有相关信息，请说明并提供你的建议。"""
        else:
            user_content = query
        
        messages.append({
            "role": "user",
            "content": user_content
        })
        
        # 调用大模型生成回复
        try:
            response = self.llm_client.chat.completions.create(
                model=self.config.QWEN_MODEL,
                messages=messages,
                temperature=0.7,
                max_tokens=1000
            )
            
            reply = response.choices[0].message.content.strip()
            return reply
            
        except Exception as e:
            print(f"生成回复错误: {e}")
            import traceback
            traceback.print_exc()
            return "抱歉，我现在无法回答您的问题。请稍后再试。"


# 全局知识抽取服务实例
rag_service = None

def get_rag_service():
    """获取知识抽取服务实例（单例模式）"""
    global rag_service
    if rag_service is None:
        rag_service = RAGService()
    return rag_service

